﻿#include "SodSystem.h"
#include "DataObject.h"
#include "Mesh.h"
#include <vtkm/VectorAnalysis.h>

#include <vtkm/worklet/WorkletMapField.h>
#include <vtkm/worklet/WorkletMapTopology.h>

#include <vtkm/exec/ParametricCoordinates.h>

struct InitialConditionOnCell : vtkm::worklet::WorkletVisitCellsWithPoints
{
  using ValueType = Vec<FloatDefault, 5>;

  using ControlSignature = void(CellSetIn, FieldInPoint, FieldOut density, FieldOut pressure);
  using ExecutionSignature = void(CellShape, _2, _3, _4);

  template<typename CellShapeTag, typename PType>
  VTKM_EXEC void operator()(const CellShapeTag& shape,
                            const PType& pts,
                            FloatDefault& density,
                            FloatDefault& pressure) const
  {
    auto p = (pts[0] + pts[1] + pts[2]) / 3;
    ValueType value = ComputeValue(p);
    density = value[0];
    pressure = value[4];
  }

  //VTKM_EXEC ValueType ComputeValue(const Vec3f& p) const
  //{
  //  ValueType primary_value;
  //  auto x = p[0];
  //  if (x < 0.5)
  //  {
  //    primary_value = ValueType{ 1, 0, 0, 0, 1 };
  //  }
  //  else
  //  {
  //    primary_value = ValueType{ 0.125, 0, 0, 0, 0.1 };
  //  }

  //  return primary_value;
  //}

  VTKM_EXEC ValueType ComputeValue(const Vec3f& p) const
  {
    ValueType primary_value;
    Vec3f center{ 3, 3, 0 };
    auto x = p[0];
    if (vtkm::Magnitude(p - center) < 1)
    {
      primary_value = ValueType{ 1, 0, 0, 0, 1 };
    }
    else
    {
      primary_value = ValueType{ 0.125, 0, 0, 0, 0.1 };
    }

    return primary_value;
  }
};

struct InitialConditionOnPoint : vtkm::worklet::WorkletMapField
{
  using ControlSignature = void(FieldIn point, FieldOut velocity);
  using ExecutionSignature = void(_1, _2);

  VTKM_EXEC
  void operator()(const Vec3f& point, Vec3f& velocity) const { velocity = Vec3f{ 0, 0, 0 }; }
};

SodSystem::SodSystem(const Configuration& cfg)
  : SGHSystem(cfg)
{
}

void SodSystem::InitialCondition()
{
  auto cellset = _data_set.GetCellSet();
  auto coords = _data_set.GetCoordinateSystem();
  auto density = GetVariableArray<Real>("density");
  auto pressure = GetVariableArray<Real>("pressure");
  auto vel = GetVariableArray<Vec3f>("velocity");

  Invoker{}(InitialConditionOnCell{}, cellset, coords, density, pressure);
  Invoker{}(InitialConditionOnPoint{}, coords, vel);
}
